package com.yanyeori.framework.log.common.wrapper;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * 将响应体中的流copy一份，实现流的重复读取
 *
 * @author chenkuan
 * @since 2022/11/18
 */
public class BufferedResponseWrapper extends HttpServletResponseWrapper {
    private final ByteArrayOutputStream baos;
    private final ServletOutputStream sos;

    public BufferedResponseWrapper(HttpServletResponse response) {
        super(response);
        baos = new ByteArrayOutputStream();
        sos = new ServletOutputStream() {
            @Override
            public void flush() throws IOException {
                super.flush();
                response.getOutputStream().flush();
                baos.flush();
            }

            @Override
            public void close() throws IOException {
                super.close();
                response.getOutputStream().close();
                baos.close();
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {

            }

            @Override
            public void write(int b) throws IOException {
                response.getOutputStream().write(b);
                baos.write(b);
            }
        };
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return sos;
    }

    public String toString() {
        return baos == null ? "" : baos.toString();
    }
}
